package 别人的教程

import (
	"fmt"
	"github.com/grafov/m3u8"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"sync"
)

/*
要使用Go语言异步下载M3U8视频，您可以结合使用goroutines和sync.WaitGroup。在下面的示例中，我们将异步下载媒体播放列表中的每个.ts片段，
并在所有片段下载完成后合并它们：
*/
func main() {
	m3u8URL := "http://example.com/path/to/video.m3u8"
	videoFilename := "output_video.ts"

	err := downloadM3U8Video(m3u8URL, videoFilename)
	if err != nil {
		fmt.Printf("Error downloading m3u8 video: %v\n", err)
	} else {
		fmt.Printf("Successfully downloaded m3u8 video to %s\n", videoFilename)
	}
}

func downloadM3U8Video(m3u8URL, videoFilename string) error {
	resp, err := http.Get(m3u8URL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	playlist, _, err := m3u8.DecodeFrom(resp.Body, true)
	if err != nil {
		return err
	}

	masterPls, ok := playlist.(*m3u8.MasterPlaylist)
	if !ok {
		return fmt.Errorf("invalid m3u8 file")
	}

	var mediaPls *m3u8.MediaPlaylist
	for _, variant := range masterPls.Variants {
		resp, err = http.Get(variant.URI)
		if err != nil {
			continue
		}
		defer resp.Body.Close()

		mediaPls, _, err = m3u8.DecodeFrom(resp.Body, true)
		if err == nil {
			break
		}
	}

	if mediaPls == nil {
		return fmt.Errorf("unable to download any media playlist")
	}

	segmentCount := len(mediaPls.Segments)
	tempFiles := make([]*os.File, segmentCount)
	var wg sync.WaitGroup

	for i, segment := range mediaPls.Segments {
		if segment == nil {
			continue
		}

		wg.Add(1)
		go func(index int, seg *m3u8.MediaSegment) {
			defer wg.Done()

			tsURL := seg.URI
			if !strings.HasPrefix(tsURL, "http") {
				tsURL = getBaseURL(m3u8URL) + tsURL
			}

			tmpFile, err := ioutil.TempFile("", "ts-segment-*.ts")
			if err != nil {
				fmt.Printf("Error creating temp file: %v\n", err)
				return
			}
			defer tmpFile.Close()

			err = downloadTS(tmpFile, tsURL)
			if err != nil {
				fmt.Printf("Error downloading ts segment from %s: %v\n", tsURL, err)
				return
			}

			tempFiles[index] = tmpFile
		}(i, segment)
	}

	wg.Wait()

	videoFile, err := os.Create(videoFilename)
	if err != nil {
		return err
	}
	defer videoFile.Close()

	for _, tmpFile := range tempFiles {
		if tmpFile == nil {
			continue
		}

		tmpFile.Seek(0, 0)
		_, err = io.Copy(videoFile, tmpFile)
		if err != nil {
			return err
		}

		os.Remove(tmpFile.Name())
	}

	return nil
}

func downloadTS(tmpFile *os.File, tsURL string) error {
	resp, err := http.Get(tsURL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	_, err = io.Copy(tmpFile, resp.Body)
	return err

}

func getBaseURL(url string) string {
	lastIndex := strings.LastIndex(url, "/")
	return url[:lastIndex+1]
}
